from machine import Pin
import usocket as socket
import network
import time
import _thread
import gc

# 配置参数
WIFI_SSID = "LAPTOP-D3RPTUAA 6358"      # WiFi名称，与ESP32-CAM相同
WIFI_PASSWORD = "19234949992.ch"    # WiFi密码，与ESP32-CAM相同
UDP_PORT = 5005                # 接收UDP图像的端口，需与ESP32-CAM设置一致
HTTP_PORT = 80                 # HTTP服务器端口

# 全局变量
latest_frame = None            # 存储最新的图像帧
frame_lock = _thread.allocate_lock()  # 线程锁，确保图像访问安全
is_connected = False           # WiFi连接状态

# WiFi连接函数
def connect_wifi(ssid, password):
    global is_connected
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    
    if not wlan.isconnected():
        print(f"连接到WiFi: {ssid}")
        wlan.connect(ssid, password)
        
        # 等待连接
        max_wait = 30
        while max_wait > 0:
            if wlan.status() < 0 or wlan.status() >= 3:
                break
            max_wait -= 1
            print("等待连接...")
            time.sleep(1)
        
        # 检查连接状态
        if wlan.status() != 3:
            print("WiFi连接失败")
            return None
        else:
            is_connected = True
            status = wlan.ifconfig()
            print(f"连接成功，IP地址: {status[0]}")
            return status[0]
    else:
        is_connected = True
        status = wlan.ifconfig()
        print(f"已连接，IP地址: {status[0]}")
        return status[0]

# UDP图像接收线程
def udp_receiver_thread():
    global latest_frame
    
    # 创建UDP套接字
    udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_socket.bind(('0.0.0.0', UDP_PORT))
    udp_socket.settimeout(1.0)  # 超时设置
    print(f"UDP服务器启动，监听端口 {UDP_PORT}")
    
    while True:
        try:
            # 接收总数据包数量
            data, addr = udp_socket.recvfrom(1024)
            if not data:
                continue
                
            try:
                total_packets = int(data.decode())
                if total_packets <= 0:
                    continue
                
                # 接收所有数据包
                packets = []
                for i in range(total_packets):
                    packet, _ = udp_socket.recvfrom(1400)  # 与发送端包大小一致
                    packets.append(packet)
                
                # 重组图像数据
                frame = b''.join(packets)
                
                # 线程安全地更新最新帧
                with frame_lock:
                    latest_frame = frame
                
                # 打印接收信息
                print(f"从 {addr} 接收图像，大小: {len(frame)} 字节，数据包数量: {total_packets}")
                
            except ValueError:
                print("无效的数据包数量")
            except Exception as e:
                print(f"接收数据包错误: {e}")
                
        except OSError as e:
            # 超时是正常的，继续等待下一组数据
            if str(e) != "ETIMEDOUT":
                print(f"UDP接收错误: {e}")
        except Exception as e:
            print(f"UDP线程错误: {e}")
        
        # 内存回收
        gc.collect()

# HTTP服务器处理函数
def handle_client(client_sock):
    global latest_frame
    try:
        # 接收客户端请求
        request = client_sock.recv(1024).decode()
        if not request:
            return
        
        # 解析请求路径
        path = request.split(' ')[1]
        
        # 主页请求
        if path == '/':
            html = """<!DOCTYPE html>
<html>
<head>
    <title>ESP32-CAM 图像流</title>
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <style>
        body { 
            font-family: Arial, sans-serif; 
            text-align: center; 
            margin: 0; 
            padding: 20px; 
            background-color: #f0f0f0;
        }
        h1 { color: #333; }
        #video-container { 
            margin: 20px auto; 
            max-width: 100%; 
            box-shadow: 0 0 10px rgba(0,0,0,0.3);
            background-color: #000;
            border-radius: 5px;
            overflow: hidden;
        }
        #stream { 
            max-width: 100%; 
            max-height: 80vh; 
        }
        .status { 
            padding: 10px; 
            margin: 10px auto; 
            border-radius: 5px; 
            max-width: 600px;
        }
        .connected { background-color: #dff0d8; color: #3c763d; }
        .disconnected { background-color: #f2dede; color: #a94442; }
    </style>
</head>
<body>
    <h1>ESP32-CAM 实时图像</h1>
    <div class="status" id="connection-status">连接状态: 等待图像...</div>
    <div id="video-container">
        <img id="stream" src="" alt="实时摄像头图像">
    </div>
    <script>
        // 更新图像的函数
        function updateImage() {
            const img = document.getElementById('stream');
            const status = document.getElementById('connection-status');
            
            // 每次添加随机数防止缓存
            img.src = "/image?" + new Date().getTime();
            
            // 图像加载成功
            img.onload = function() {
                status.className = "status connected";
                status.textContent = "连接状态: 正常 (" + new Date().toLocaleTimeString() + ")";
            };
            
            // 图像加载失败
            img.onerror = function() {
                status.className = "status disconnected";
                status.textContent = "连接状态: 图像获取失败 (" + new Date().toLocaleTimeString() + ")";
            };
        }
        
        // 初始加载和定时更新 (每100ms更新一次)
        updateImage();
        setInterval(updateImage, 100);
    </script>
</body>
</html>"""
            response = f"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: {len(html)}\r\n\r\n{html}"
            client_sock.send(response.encode())
        
        # 图像数据请求
        elif path.startswith('/image'):
            with frame_lock:
                if latest_frame:
                    # 发送JPEG图像
                    response = f"HTTP/1.1 200 OK\r\nContent-Type: image/jpeg\r\nContent-Length: {len(latest_frame)}\r\n\r\n"
                    client_sock.send(response.encode() + latest_frame)
                else:
                    # 没有图像数据时发送404
                    client_sock.send("HTTP/1.1 404 Not Found\r\n\r\nNo image available".encode())
        
        # 其他路径
        else:
            client_sock.send("HTTP/1.1 404 Not Found\r\n\r\nPage not found".encode())
            
    except Exception as e:
        print(f"处理客户端错误: {e}")
    finally:
        client_sock.close()

# HTTP服务器线程
def http_server_thread():
    # 创建TCP套接字
    tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    tcp_socket.bind(('0.0.0.0', HTTP_PORT))
    tcp_socket.listen(5)
    print(f"HTTP服务器启动，端口 {HTTP_PORT}")
    
    while True:
        try:
            client_sock, client_addr = tcp_socket.accept()
            print(f"新客户端连接: {client_addr}")
            handle_client(client_sock)
        except Exception as e:
            print(f"HTTP服务器错误: {e}")
        
        # 内存回收
        gc.collect()

# 主函数
def main():
    # 连接WiFi
    ip_addr = connect_wifi(WIFI_SSID, WIFI_PASSWORD)
    if not ip_addr:
        print("无法连接WiFi，程序退出")
        return
    
    print(f"请在浏览器中访问: http://{ip_addr}")
    
    # 启动UDP接收线程
    _thread.start_new_thread(udp_receiver_thread, ())
    
    # 启动HTTP服务器 (在主线程中运行)
    http_server_thread()

# 启动程序
if __name__ == "__main__":
    main()
